#include "ParseBaseNodes.h"
#include "NodeVisitor.h"
#include "TypeCheckDispatcher.h"

using namespace std;

void TypeCheckVisitor::visit(ConstructorDeclarationNode*n) {
	visit((MethodDeclarationNode*)n);
	//Make sure a default declaration is available, since JOOS has no explicit super calls
	ClassDeclarationNode *super = TypeCheckGetContainingClass(n)->super;
	if (super) {
		//Build a temp typed node that contains the signature
		TypedNode tn;
		tn.identifier = super->identifier + "$";

		//Search on the signature
		set<TypedNode*, compare_decls>::iterator constructor_declaration = super->elements_in_scope[DeclarationNode::Method].find(&tn);

		//Error if the constructor does not exist;
		error |= constructor_declaration == super->elements_in_scope[DeclarationNode::Method].end() || !CanAccessMember(*constructor_declaration, TypeCheckGetContainingClass(n), TypeCheckGetContainingClass(n));

	}
}
void TypeCheckVisitor::visit(FieldDeclarationNode*n) {
	if ((*n)["Expression"]) visit((LocalVariableDeclarationNode*)n);
	visit((MethodDeclarationNode*)n);
}


void TypeCheckVisitor::visit(IfStatementNode *) {
	while (!conditional_metas.empty()) {
		error |= !(conditional_metas.top().category & TypeMeta::Boolean);
		conditional_metas.pop();
	}
}
void TypeCheckVisitor::visit(IfThenElseStatementNode *) {
	while (!conditional_metas.empty()) {
		error |= !(conditional_metas.top().category & TypeMeta::Boolean);
		conditional_metas.pop();
	}
}
void TypeCheckVisitor::visit(ForStatementNode *) {
	while (!conditional_metas.empty()) {
		error |= !(conditional_metas.top().category & TypeMeta::Boolean);
		conditional_metas.pop();
	}
}
void TypeCheckVisitor::visit(WhileStatementNode *) {
	while (!conditional_metas.empty()) {
		error |= !(conditional_metas.top().category & TypeMeta::Boolean);
		conditional_metas.pop();
	}
}

void TypeCheckVisitor::visit(ExpressionNode *n) {
	string parent_type = n->parent->Type();
	if (parent_type == "IfStatement" || parent_type == "IfThenElseStatement" || parent_type == "WhileStatement" || parent_type == "ForStatement") conditional_metas.push(metas.top());
}

void TypeCheckVisitor::visit(ConditionalOrExpressionNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1 = metas.top();
	metas.pop();

	TypeMetaReference result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::ConditionalOr, t1, t2);
	error |= result.category == TypeMeta::Error;

	
	if (result.boolean_literal) {
		n->ReplaceMeWith(result.boolean_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(ConditionalAndExpressionNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1 = metas.top();
	metas.pop();

	TypeMetaReference result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::ConditionalAnd, t1, t2);
	error |= result.category == TypeMeta::Error;

	
	if (result.boolean_literal) {
		n->ReplaceMeWith(result.boolean_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(InclusiveOrExpressionNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1 = metas.top();
	metas.pop();

	TypeMetaReference result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::InclusiveOr, t1, t2);
	error |= result.category == TypeMeta::Error;

	if (result.boolean_literal) {
		n->ReplaceMeWith(result.boolean_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(ExclusiveOrExpressionNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1 = metas.top();
	metas.pop();

	TypeMetaReference result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::ExclusiveOr, t1, t2);
	error |= result.category == TypeMeta::Error;

	if (result.boolean_literal) {
		n->ReplaceMeWith(result.boolean_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(AndExpressionNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1 = metas.top();
	metas.pop();

	TypeMetaReference result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::And, t1, t2);
	error |= result.category == TypeMeta::Error;

	
	if (result.boolean_literal) {
		n->ReplaceMeWith(result.boolean_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(EqualityExpressionNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1 = metas.top();
	metas.pop();

	TypeMetaReference result = TypeMetaReference();
	if ((*n)["=="]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::EqualTo, t1, t2);
	else if ((*n)["!="]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::NotEqualTo, t1, t2);

	error |= result.category == TypeMeta::Error;
	
	if (result.boolean_literal) {
		n->ReplaceMeWith(result.boolean_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(RelationalExpressionNode *n) {
	if (n->children.size() == 1) return;

	TypeMetaReference t1, t2, result = TypeMetaReference();
	if ((*n)["instanceof"]) {
		t1 = metas.top();
		metas.pop();
		ReferenceTypeNode *r = ((ReferenceTypeNode*)(*n)["ReferenceType"]);
		t2 = TypeMetaReference(r->type_meta.category, r->type_meta.reference, 0x0);
		result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::InstanceOf, t1, t2);
	} else {

		t2 = metas.top();
		metas.pop();
		t1 = metas.top();
		metas.pop();

		if ((*n)["<"]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::LessThan, t1, t2);
		else if ((*n)[">"]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::GreaterThan, t1, t2);
		else if ((*n)["<="]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::LessThanEqualTo, t1, t2);
		else if ((*n)[">="]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::GreaterThanEqualTo, t1, t2);

	}
	error |= result.category == TypeMeta::Error;
	
	if (result.boolean_literal) {
		n->ReplaceMeWith(result.boolean_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(AdditiveExpressionNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1 = metas.top();
	metas.pop();

	TypeMetaReference result = TypeMetaReference();
	if ((*n)["+"]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Addition, t1, t2);
	else if ((*n)["-"]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Subtraction, t1, t2);

	error |= result.category == TypeMeta::Error;
	
	if (result.integer_literal) {
		n->ReplaceMeWith(result.integer_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(MultiplicativeExpressionNode *n) {
	if (n->children.size() == 1) return;

	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1 = metas.top();
	metas.pop();

	TypeMetaReference result = TypeMetaReference();
	if ((*n)["*"]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Multiplication, t1, t2);
	else if ((*n)["/"]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Division, t1, t2);
	else if ((*n)["%"]) result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Modulus, t1, t2);

	error |= result.category == TypeMeta::Error;
	
	if (result.integer_literal) {
		n->ReplaceMeWith(result.integer_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(UnaryExpressionNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t1 = metas.top();
	metas.pop();


	TypeMetaReference result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Negation, t1, t1);
	error |= result.category == TypeMeta::Error;

	
	if (result.integer_literal) {
		n->ReplaceMeWith(result.integer_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(UnaryExpressionNotPlusMinusNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t1 = metas.top();
	metas.pop();

	TypeMetaReference result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Complement, t1, t1);
	error |= result.category == TypeMeta::Error;

	
	if (result.boolean_literal) {
		n->ReplaceMeWith(result.boolean_literal);
	} else metas.push(result);
}
void TypeCheckVisitor::visit(LocalVariableDeclarationNode *n) {
	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1(n->type_meta.category, n->type_meta.reference, 0x0);

	metas.push(ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Assignment, t1, t2));
	error |= metas.top().category == TypeMeta::Error;
}
void TypeCheckVisitor::visit(AssignmentNode *n){
	if (n->children.size() == 1) return;

	TypeMetaReference t2 = metas.top();
	metas.pop();
	TypeMetaReference t1 = metas.top();
	metas.pop();

	metas.push(ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Assignment, t1, t2));
	error |= metas.top().category == TypeMeta::Error;
}
void TypeCheckVisitor::visit(ReturnStatementNode *n) {
	//Void return statement
	if (n->children.size() == 2) {
		ParentNode *p = n->parent;
		for (; p->Type() != "MethodDeclaration" && p->Type() != "ConstructorDeclaration"; p = p->parent);
		error |= !(((TypedNode*)p)->type_meta.category & TypeMeta::Void);

	//Return statement has an expression
	} else {
		//Get the type of the return value off the top of the stack
		TypeMetaReference return_value_type = metas.top();
		metas.pop();

		//Find the method we're returning from
		ParentNode *p = n->parent;
		for (; p->Type() != "MethodDeclaration" && p->Type() != "ConstructorDeclaration"; p = p->parent);

		//Push the method return value type onto the stack, then the return value's type
		TypedNode *m = (TypedNode*)p;
		metas.push(TypeMetaReference(m->type_meta.category, m->type_meta.reference, 0x0));
		metas.push(return_value_type);

		//Stack imitates what the assignment node wants (IE type = type), because to the return value must be assignable to the return type
		visit((AssignmentNode*)n);
	}
}
void TypeCheckVisitor::visit(CastExpressionNode *n) {
	TypeMetaReference result = ITypeCheckDispatcher::ResolveType(ITypeCheckDispatcher::Casting, TypeMetaReference(n->type_meta.category, n->type_meta.reference, 0x0), metas.top()); 
	metas.pop();
	error |= result.category == TypeMeta::Error;
	
	if (result.integer_literal) {
		result.integer_literal->underlying_type = result.category;
		n->ReplaceMeWith(result.integer_literal);
	} else metas.push(result);
}

void TypeCheckVisitor::visit(PrimaryNoNewArrayNode *n) {
	//The "this" node was here, but was deleted, because fuck that guy. But we still need to handle the case that he was here.
	if (n->children.empty()) {

		ParentNode *p = n->parent;
		for (; p->Type() != "FieldDeclaration" && p->Type() != "MethodDeclaration" && p->Type() != "ConstructorDeclaration"; p = p->parent);

		//Can't access 'this' from a static context
		error |= (((ModifiedNode*)p)->modifier & ModifiedNode::Static) != 0x0;

		//For "this" push on the class that the "this" is in
		metas.push(TypeMetaReference(TypeMeta::Reference, TypeCheckGetContainingClass(n), 0x0));
	}
}
void TypeCheckVisitor::visit(ArrayCreationExpressionNode *n) {
	//The expression must be an integral type
	error |= !(metas.top().category & (TypeMeta::Int | TypeMeta::Short | TypeMeta::Byte | TypeMeta::Char));
	if (error) return;
	metas.pop();

	//Push in this new array type
	metas.push(TypeMetaReference((TYPE)(n->type_meta.category | TypeMeta::Array), n->type_meta.reference, 0x0));
}
void TypeCheckVisitor::visit(ClassInstanceCreationExpressionNode *n) {
	error |= n->type_meta.reference->Type() == "InterfaceDeclaration" || (n->type_meta.reference->modifier & ModifiedNode::Abstract) != 0x0;
	if (error) return;

	//metas need to be pop'd in reverse order. Arugment metas are actually first. Build the method signature, before determining method context
	string argument_signature = BuildArgumentSignature(n);

	//Build a temp typed node that contains the signature
	TypedNode tn;
	tn.identifier = n->type_name + "$" + argument_signature;

	//Search on the signature
	set<TypedNode*, compare_decls>::iterator constructor_declaration = n->type_meta.reference->elements_in_scope[DeclarationNode::Method].find(&tn);

	//Error if the constructor does not exist;
	error |= constructor_declaration == n->type_meta.reference->elements_in_scope[DeclarationNode::Method].end() || !CanAccessConstructor(*constructor_declaration, TypeCheckGetContainingClass(n));
	if (error) return;

	//Store the pointer to this constructor, if we're actually the creation expression
	n->constructor = (ConstructorDeclarationNode*)*constructor_declaration;

	//Push in the meta stating that a new object with this class was created
	metas.push(TypeMetaReference(n->type_meta.category, n->type_meta.reference, 0x0));
}
void TypeCheckVisitor::visit(FieldAccessNode *n) {
	TypeMetaReference object_type = metas.top();
	metas.pop();

	string field_name = ((LeafNode*)(*n)["identifier"])->lexeme;

	if (object_type.category & TypeMeta::Array) {
		//Must be the field name length and we can't be assigning to it
		error |= field_name != "length" || n->parent->Type() == "LeftHandSide";
		n->field = 0x0;
		metas.push(TypeMetaReference(TypeMeta::Int, 0x0, 0x0));
		return;
	}

	//Error if this is a primitive
	error |= !(object_type.category & (TypeMeta::Reference | TypeMeta::String));
	if (error) return;

	bool field_is_static = false;

	//Put the field name into a lookup node
	TypedNode tn;
	tn.identifier = field_name;

	
	EnvironmentNode *env = object_type.reference;
	set<TypedNode*, compare_decls>::iterator field_declaration;
	for (; env; env = env->next_environment) {
		field_declaration = env->elements_in_scope[DeclarationNode::Expression].find(&tn);
		if (field_declaration != env->elements_in_scope[DeclarationNode::Expression].end()) break;
	}
	

	//Error if we ran out of environments to check or we can't access this field
	error |= !env || !CanAccessMember(*field_declaration, TypeCheckGetContainingClass(n), object_type.reference);
	if (error) return;

	if (field_is_static) {
		error |= !((**field_declaration).modifier & TypedNode::Static);
	}

	n->field = (FieldDeclarationNode*)*field_declaration;

	metas.push(TypeMetaReference((**field_declaration).type_meta.category, (**field_declaration).type_meta.reference, 0x0));
}
void TypeCheckVisitor::visit(MethodInvocationNode *n) {

	//metas need to be pop'd in reverse order. Arugment metas are actually first. Build the method signature, before determining method context
	string argument_signature = BuildArgumentSignature(n);

	NameNode *nn = (NameNode*)(*n)["Name"];

	bool method_must_be_static = false, implicit_this = false;
	//We found that the invocation was from a name (opposed to a subexpression)
	if (nn) {
		//The name, up to the method identifier, could be only a class reference. This means the method needs to be static
		method_must_be_static = nn->declaration_is_static_context;
		implicit_this = nn->children.empty();
	}

	TypeMetaReference method_class_type = metas.top();
	metas.pop();

	//Can't call a method on a primitive (this is not where the array.length call is handled)
	error |= !(method_class_type.category & (TypeMeta::Reference | TypeMeta::String));
	if (error) return;

	//Build a temp typed node that contains the signature
	TypedNode tn;
	tn.identifier = n->method_identifier + argument_signature;

	//Search on the signature in the recognized type
	set<TypedNode*, compare_decls>::iterator method_declaration = method_class_type.reference->elements_in_scope[DeclarationNode::Method].find(&tn);

	//Error if the method does not exist or we don't have access to it
	error |= (method_declaration == method_class_type.reference->elements_in_scope[DeclarationNode::Method].end() || !CanAccessMember(*method_declaration, TypeCheckGetContainingClass(n), (method_must_be_static? TypeCheckGetContainingClass(n) : method_class_type.reference)));
	if (error) return;

	bool method_is_static = ((**method_declaration).modifier & TypedNode::Static) != 0x0;
	error |= method_is_static != method_must_be_static;

	//Can't use implicit this on static calls
	if (method_is_static) {
		error |= implicit_this;
	}

	//Store the pointer to this method in the invocation
	n->method = (MethodDeclarationNode*)*method_declaration;

	//Push in the meta stating that this method returns a value of this type
	metas.push(TypeMetaReference((**method_declaration).type_meta.category, (**method_declaration).type_meta.reference, 0x0));

}
void TypeCheckVisitor::visit(ArrayAccessNode *n) {
	//That index must be an integer type
	error |= !(metas.top().category & (TypeMeta::Int | TypeMeta::Short | TypeMeta::Byte | TypeMeta::Char));
	metas.pop();

	TYPE &cat = metas.top().category;
	//Make sure the type that we're accessing is an array
	error |= !(cat & TypeMeta::Array);

	//Whatever the array was, it's just the component type now
	cat = (TYPE)(cat ^ TypeMeta::Array);

	n->array_type = metas.top();
}

void TypeCheckVisitor::visit(NameNode *n) {

	metas.push(TypeMetaReference(n->declaration_reference->type_meta.category, n->declaration_reference->type_meta.reference, 0x0));
}
void TypeCheckVisitor::visit(IntegerLiteralNode *n) {
	metas.push(TypeMetaReference(n->underlying_type, 0x0, n));
}
void TypeCheckVisitor::visit(CharacterLiteralNode *n) {
	metas.push(TypeMetaReference(TypeMeta::Char, 0x0, n));
}
void TypeCheckVisitor::visit(StringLiteralNode *n) {
	metas.push(TypeMetaReference(TypeMeta::String, TypeMeta::string_declaration, n));
}
void TypeCheckVisitor::visit(BooleanLiteralNode *n) {
	metas.push(TypeMetaReference(TypeMeta::Boolean, 0x0, n));
}
void TypeCheckVisitor::visit(NullLiteralNode *n) {
	metas.push(TypeMetaReference(TypeMeta::Null, 0x0, n));
}



string TypeCheckVisitor::BuildArgumentSignature(ParentNode *n) {

	unsigned int num_arguments = 0;
	for (vector<SP<ParseNode> >::iterator it = n->children.begin(); it < n->children.end(); ++it) {
		if ((**it).Type() == "Expression") ++num_arguments; 
	}

	string argument_signature = string();

	//Remember that metas are pop'd in reverse order. Last arguments are pop'd first
	//Thus why the inserts at 0 (IE. prepending)
	while (num_arguments-- > 0) {
		TypeMetaReference t = metas.top();
		metas.pop();
		if (!t.reference) {
			string mods = (t.category & TypeMeta::Array ? "$" : "");

			switch (t.category & ~TypeMeta::Array) {
				case TypeMeta::Boolean:
					argument_signature.insert(0, "$boolean" + mods);
					break;
				case TypeMeta::Byte:
					argument_signature.insert(0, "$byte" + mods);
					break;
				case TypeMeta::Short:
					argument_signature.insert(0, "$short" + mods);
					break;
				case TypeMeta::Int:
					argument_signature.insert(0, "$int" + mods);
					break;
				case TypeMeta::Char:
					argument_signature.insert(0, "$char" + mods);
					break;
				default:
					error = true;
					return string();
			}

		} else argument_signature.insert(0, '$' + t.reference->package_name + t.reference->identifier + (t.category & TypeMeta::Array ? "$" : ""));
	}

	return argument_signature;
}